from langchain.agents import create_pandas_dataframe_agent
from langchain.chat_models import ChatOpenAI
from langchain.agents.agent_types import AgentType

from langchain.llms import OpenAI
import pandas as pd

# 读取数据
df = pd.read_csv("titanic.csv")

# 使用 ZERO_SHOT_REACT_DESCRIPTION 代理类型进行初始化
agent = create_pandas_dataframe_agent(OpenAI(temperature=0), df, verbose=True)

# 使用OpenAI函数，使用 OPENAI_FUNCTIONS 代理类型进行初始化
agent = create_pandas_dataframe_agent(
    ChatOpenAI(temperature=0, model="gpt-3.5-turbo-0613"),
    df,
    verbose=True,
    agent_type=AgentType.OPENAI_FUNCTIONS,
)

# 执行问答
print(agent.run("有多少行？"))